// ripeasm.c
// Implementation of the RipeScript assembler.
//
// Copyright (C) 2008  Maksim Sipos <msipos@mailc.net>
//
// This file is a part of the Ripe language project.
//
// Ripe is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

#include <cstdlib>
#include <cstdio>
#include <fstream>
#include <getopt.h>
#include "util/standard.hpp"
#include "util/error.hpp"
#include "util/lex.hpp"
#include "util/token.hpp"
#include "util/packrat.hpp"
#include "util/serialize.hpp"
#include "adts/intern.hpp"
#include "vm/vmunit.hpp"
#include "vm/vmimage.hpp"

using namespace Ripe;
using std::ifstream;
using std::ofstream;
using std::printf;
using std::atoi;

// Assembler token
class AToken : public Token {
  public:
    // Used if AToken is a TOK_ID, TOK_STRING or TOK_ADDRESS. Pointer to an
    // interned string.
    char* text;
    
    // Used if AToken is a TOK_INTEGER or TOK_REGISTER
    int32_t integer;
};

// Assembler token types:
#define TOK_COMMENT   0
#define TOK_WHITE     1
#define TOK_NEWLINE   2
#define TOK_ID        3
#define TOK_COMMA     4
#define TOK_STRING    5
#define TOK_AMPERSAND 6
#define TOK_ADDRESS   7 // Specially constructed by the parser, out of a
                        // '&' IDENTIFIER sequence
#define TOK_COLON     8
#define TOK_INTEGER   9
#define TOK_GLOBALID 10
#define TOK_REGISTER 11

#define TOK_CALL     21
#define TOK_ASSIGN   22
#define TOK_CALLR    23
#define TOK_ACCEPT   24
#define TOK_JUMP     25
#define TOK_JUMPIF   26
#define TOK_RET      27

// Grammar non-terminals:
#define PROGRAM   1
#define STATEMENT 2
#define ARG       3

// Special semantic values
#define EMPTY_STATEMENT 1

// Abstract syntax tree:
// One statement in the assembly source
class Statement {
  public:
    char* label;
    uint32_t instruction;
    Vector<AToken*> args;
};

// The whole source file
class Program {
  public:
    Vector<Statement*> statements;
};

// General error while trying to assemble a file
class ErrorAssemble : public Exception {
};

// A structure that combines the symbol table, internment table, direct-
// value tables, etc.
class Table {
  public:
    Table();
  
    Intern intern;
    
    uint32_t n_symbols;
    Map<char*, uint32_t> symbols;
    uint32_t query_symbol(char* s);

    uint32_t n_directs;
    Map<char*, uint32_t> directs_strings;
    Map<int32_t, uint32_t> directs_integers;
    uint32_t query_string(char* s);
    uint32_t query_integer(int32_t integer);
};

Table::Table()
{
  n_symbols = 0;
  n_directs = 0;
}

uint32_t Table::query_symbol(char* s)
{
  uint32_t n;
  if (not symbols.get(s, &n)){
    n = n_symbols;
    symbols.set(s, n);
    n_symbols++;
  }
  return n;
}

uint32_t Table::query_string(char* s)
{
  uint32_t n;
  if (not directs_strings.get(s, &n)){
    n = n_directs;
    directs_strings.set(s, n);
    n_directs++;
  }
  return n;
}

uint32_t Table::query_integer(int32_t integer)
{
  uint32_t n;
  if (not directs_integers.get(integer, &n)){
    n = n_directs;
    directs_integers.set(integer, n);
    n_directs++;
  }
  return n;
}

// Display help message.
static void help_message(char* arg1)
{
  printf("Usage: %s [OPTIONS] FILES\n"
         "Assemble the given source files into a VM binary.\n"
         "\n"
         "Possible options:\n"
         "  -h, --help         Display this message.\n"
         "  -o, --output FILE  Set the output filename.\n"
         "\n"
         "Report bugs to <msipos@mailc.net>\n", arg1);
}

// Returns true if a given token is an instruction token
static bool token_is_instruction(Token* token)
{
  switch(token->type){
    case TOK_CALL:
    case TOK_CALLR:
    case TOK_ASSIGN:
    case TOK_ACCEPT:
    case TOK_JUMP:
    case TOK_JUMPIF:
    case TOK_RET:
      return true;
    default:
      return false;
  }
}

// Returns true if a given token is a valid instruction argument
static inline bool token_is_arg(Token* token)
{
  switch(token->type){
    case TOK_ID:
    case TOK_STRING:
    case TOK_INTEGER:
    case TOK_GLOBALID:
    case TOK_REGISTER:
      return true;
    default:
      return false;
  }
}

// Find the VM number of the given instruction's token.
static uint32_t token_to_instruction(int32_t token_type)
{
  // TODO: Best to make a table of tokens to handle this and above
  // token_is_instruction and token_is_arg.
  switch(token_type){
    case TOK_CALL:
      return VM_CALL;
    case TOK_CALLR:
      return VM_CALLR;
    case TOK_ASSIGN:
      return VM_ASSIGN;
    case TOK_JUMP:
      return VM_JUMP;
    case TOK_JUMPIF:
      return VM_JUMPIF;
    case TOK_RET:
      return VM_RET;
    case TOK_ACCEPT:
      return VM_ACCEPT;
    default:
      assert_never();
  }
}

// program ::= stmt *
static Semantic match_program(Parsing* p, int precedence)
{
  Semantic stmt;
  Program* program = new Program;

  while ((stmt = p->match(STATEMENT, 0)) != NULL) {
    if (uint32_t(stmt) != EMPTY_STATEMENT){
      program->statements.append((Statement*) stmt);
    }
  }
  return (Semantic) program;
}

// stmt ::= [ID ':'] INSTRUCTION [arg (, arg)*] '\n'
static Semantic match_statement(Parsing* p, int precedence)
{
  Token* token;
  char* label = NULL;

  // Check if there is a label.
  if (p->consume(TOK_ID, &token)) {
    Token* token2;
    if (p->consume(TOK_COLON, &token2)){
      label = ((AToken*) token)->text;
    } else {
      return NULL;
    }
  }
    
  if (p->read(&token)){
    if (token_is_instruction(token)) {
      Statement* stmt = rnew Statement;
      stmt->label = label;
      stmt->instruction = token_to_instruction(token->type);

      Token* arg;
      // Try to read first arg
      if ((arg = (Token*) p->match(ARG, 0))) {
        // Yes, first argument read.
        stmt->args.append((AToken*) arg);

        for(;;){
          if (p->consume(TOK_COMMA, &token)) {
            if ((arg = (AToken*) p->match(ARG, 0))) {
              stmt->args.append((AToken*) arg);
            } else {
              return NULL;
            }
          } else if (p->consume(TOK_NEWLINE, &token)) {
            return (Semantic) stmt;
          } else {
            return NULL;
          }
        }
      } else if (p->consume(TOK_NEWLINE, &token)) {
        return (Semantic) stmt;
      } else {
        return NULL;
      }
    } else if (token->type == TOK_NEWLINE) {
      return (Semantic) EMPTY_STATEMENT;
    }
  }
  return NULL;
}

// arg ::= '&' ID | ID | STRING
// Semantic value is the token matched (changed into TOK_ADDRESS in the first
// case).
static Semantic match_arg(Parsing* p, int precedence)
{
  Token* token;
  if (p->consume(TOK_AMPERSAND, &token)){
    if (p->consume(TOK_ID, &token)){
      // '&' ID
      token->type = TOK_ADDRESS;
      return (Semantic) token;
    } else {
      // Failed matching (something else follows '&').
      return NULL;
    }
  } else if (p->read(&token) and token_is_arg(token)) {
    return (Semantic) token;
  } else {
    return NULL;
  }
}

static void assemble(Table *table, Lexer *lexer, Parsing* parsing, 
                     VMImage* vmimage, char* filename)
{
  // Tokenize the assembler file
  int tok;
  Vector<AToken*> tokens;

  while ((tok = lexer->read_token()) != ID_EOF){
    if (tok == TOK_COMMENT or tok == TOK_WHITE) {
      // Ignore the token then.
      continue;
    }
    AToken* atoken = new AToken;
    atoken->type = tok;

    // Intern up identifiers and strings. Parse integers.
    switch(tok){
      case TOK_ID:
      case TOK_STRING:
        atoken->text = table->intern.intern(lexer->token);
        break;
      case TOK_GLOBALID:
        atoken->text = table->intern.intern(lexer->token + 1);
        break;
      case TOK_INTEGER:
        atoken->integer = atoi(lexer->token);
        break;
      case TOK_REGISTER:
        atoken->integer = atoi(lexer->token + 1);
        break;
      default: ;
    }
    tokens.append(atoken);
  }

  // parse method call resets the parser.
  parsing->tokens = (Vector<Token*>*) (&tokens);
  Program* ast = (Program*) parsing->parse(PROGRAM);

  if (ast == NULL){
    ErrorAssemble e;
    e.message = "Syntax error.";
    throw e;
  }

  // Construct a map that will help resolve the labels into actual addresses
  // in the bytecode.
  Map<char*, uint32_t> labels;
  uint32_t current = 0;
  for (uint32_t i = 0; i < ast->statements.size; i++) {
    char* label = ast->statements[i]->label;
    if (label != NULL) {
      labels.set(label, current);
    }
    current += 1 + ast->statements[i]->args.size;
  }

  // Construct a code section.
  VMSectionCode* section = rnew VMSectionCode;
  section->name = gc_strdup(filename);
  vmimage->sections.append(section);
  
  // Construct the units out of the AST.
  for (uint32_t i = 0; i < ast->statements.size; i++) {
    // Handle the instruction
    Statement* stmt = ast->statements[i];
    section->units.append(unit_instruction(stmt->instruction,
                                           stmt->args.size));

    // Handle the instruction arguments
    for (uint32_t j = 0; j < stmt->args.size; j++) {
      AToken* token = stmt->args[j];
      uint32_t label_address;
      // Map an AToken argument into a Unit.
      switch(token->type){
        case TOK_ID:
          section->units.append(
            unit_symbol(table->query_symbol(token->text)));
          break;
        case TOK_STRING:
          section->units.append(
            unit_direct_value(table->query_string(token->text)));
          break;
        case TOK_INTEGER:
          section->units.append(
            unit_direct_value(table->query_integer(token->integer)));
          break;
        case TOK_GLOBALID:
          section->units.append(
            unit_global_symbol(table->query_symbol(token->text)));
          break;
        case TOK_REGISTER:
          section->units.append(unit_register(token->integer));
          break;
        case TOK_ADDRESS:
          if (not labels.get(token->text, &label_address)) {
            ErrorAssemble e;
            e.message.format("Unknown label '%s'.", token->text);
            throw e;
          }
          section->units.append(unit_address(label_address));
          break;
        default:
          assert_never();
      }      
    }
  }
}

int main(int argc, char** argv)
{
  const char* o_filename = "r.out";

  // Initialize the error-reporting system.
  error_init(argv[0]);

  struct option long_options[] =
    {
      {"help", no_argument, 0, 'h'},
      {"output", required_argument, 0, 'o'},
      {0, 0, 0, 0}
    };
  // getopt_long stores the option index here.
  int option_index = 0;

  while(true){
    int c = getopt_long(argc, argv, "ho:", long_options, &option_index);

    // End of long/short options, remainder of the arguments should be
    // interpreted like input files.
    if (c == EOF){
      break;
    }

    switch(c){
      case 'h':
        help_message(argv[0]);
        return 0;
      case 'o':
        if (optarg == NULL){
          error("Option '-o' requires an argument.");
          return 1;
        }
        o_filename = gc_strdup(optarg);
        break;
      default:
        error("Parsing command-line arguments failed.");
        return 1;
    }
  }

  if (optind >= argc){
    error("No input files specified.");
    return 1;
  }

  // Construct the assembler lexer
  Lexer lexer(3);
  lexer.add_rule(0, ENV_STAY, "[ \t]+", TOK_WHITE);
  lexer.add_rule(0, ENV_STAY, "\n", TOK_NEWLINE);
  lexer.add_rule(0, ENV_STAY, "[A-Za-z_][A-Za-z0-9_]*", TOK_ID);
  lexer.add_rule(0, ENV_STAY, "$[A-Za-z_][A-Za-z0-9_]*", TOK_GLOBALID);
  lexer.add_rule(0, ENV_STAY, ",", TOK_COMMA);
  lexer.add_rule(0, ENV_STAY, "&", TOK_AMPERSAND);
  lexer.add_rule(0, ENV_STAY, ":", TOK_COLON);
  lexer.add_rule(0, ENV_STAY, "-?[0-9]+", TOK_INTEGER);
  lexer.add_rule(0, ENV_STAY, "@[0-9]+", TOK_REGISTER);

  // Deal with strings
  lexer.add_rule(0,        1, "\"", ID_CONSUME);
  lexer.add_rule(1, ENV_STAY, "[^\"]+", TOK_STRING);
  lexer.add_rule(1,  ENV_POP, "\"", ID_CONSUME);

  // Deal with comments
  lexer.add_rule(0,        2, "#", ID_CONSUME);
  lexer.add_rule(2,  ENV_POP, "[^\n]+", TOK_COMMENT);

  // Reserved keywords (instructions)
  // Rules defined later get preference over identifiers when matching
  lexer.add_rule(0, ENV_STAY, "CALL", TOK_CALL);
  lexer.add_rule(0, ENV_STAY, "ASSIGN", TOK_ASSIGN);
  lexer.add_rule(0, ENV_STAY, "CALLR", TOK_CALLR);
  lexer.add_rule(0, ENV_STAY, "ACCEPT", TOK_ACCEPT);
  lexer.add_rule(0, ENV_STAY, "JUMP", TOK_JUMP);
  lexer.add_rule(0, ENV_STAY, "JUMPIF", TOK_JUMPIF);
  lexer.add_rule(0, ENV_STAY, "RET", TOK_RET);

  // Construct the internment/symbol/direct value table
  Table table;

  // Construct the parser.
  Parsing parsing;
  parsing.grammar.set(PROGRAM, match_program);
  parsing.grammar.set(STATEMENT, match_statement);
  parsing.grammar.set(ARG, match_arg);

  // Construct the VMImage that will be populated by assemble() function.
  VMImage vmimage;

  int arg = optind;
  while (arg < argc){
    ifstream f(argv[arg], ifstream::in | ifstream::binary);
    if (f.fail()) {
      error("Failed opening file '%s'.", argv[arg]);
      return 1;
    }
    
    try {
      lexer.set_input(&f);
      assemble(&table, &lexer, &parsing, &vmimage, argv[arg]);
    } catch (const ErrorAssemble& e) {
      error("While processing file '%s': %s", argv[arg], e.what());
      return 1;
    }
    arg++;
  }

  // Construct the VM symbols sections.
  VMSectionSymbol symbols_section;
  symbols_section.name = "symbols";
  symbols_section.symbols.set_size(table.n_symbols);  
  {
    MapIterator<char*, uint32_t> iter(&(table.symbols));
    char* key; uint32_t value;
    while (iter.get(&key, &value)){
      symbols_section.symbols.set(value, key);
    }
  }
  vmimage.sections.append(&symbols_section);

  // Construct the VM directs sections.
  VMSectionDirect directs_section;
  directs_section.name = "directs";
  directs_section.directs.set_size(table.n_directs);
  DirectValue dv;
  {
    MapIterator<char*, uint32_t> iter(&table.directs_strings);
    char* key; uint32_t value;
    while (iter.get(&key, &value)){
      dv.type = DV_STRING;
      dv.string = key;
      directs_section.directs.set(value, dv);
    }
  }
  { 
    MapIterator<int32_t, uint32_t> iter(&table.directs_integers);
    int32_t key; uint32_t value;
    while (iter.get(&key, &value)){
      dv.type = DV_INTEGER;
      dv.integer = key;
      directs_section.directs.set(value, dv);
    }
  }
  vmimage.sections.append(&directs_section);

  // Serialize the VMImage.
  ofstream f(o_filename, ofstream::out | ofstream::binary);
  OSerializer serializer(&f);
  vmimage.save(&serializer);
  
  return 0;
}
